import mujoco
import os
import numpy as np
import dearpygui.dearpygui as dpg
import mujoco
import numpy as np
import glfw
import time
import mujoco.viewer
## feidedaoRobot 飞的岛

current_directory = os.getcwd()
parent_directory = os.path.dirname(current_directory)
print("路径：", current_directory)

modelPath = current_directory+'/' # parent_directory + "/fdd_robot_control/"
xml_FILENAME = "robot_description/mjcf/robot_scene.xml"

# 加载模型
xml_path = modelPath+xml_FILENAME
m = mujoco.MjModel.from_xml_path(xml_path)
d = mujoco.MjData(m)

targetKp = np.array([300, 300, 500, 500, 500, 300,
                     300, 300, 500, 500, 500, 300])
targetPos = np.array([0, 0, -0.1, 0.1, 0, 0,
                      0, 0,  0.1, -0.1, 0, 0])
posDef2 = np.array([0, 0, 0, 0, 0, 0,
                    0, 0, 0, 0, 0, 0])

targetKd = 3
targetVel = 0
print(m.opt.timestep)  # 运行周期是 0.002s

countPos_ = 0
# # 初始化 Dear PyGui
# dpg.create_context()
# dpg.create_viewport(title='MuJoCo with UI', width=800, height=600)
#
# # 创建一个窗口
# with dpg.window(label="Simulation Control", width=200, height=100):
#     # 添加一个按钮
#     def button_callback(sender, app_data, user_data):
#         global countPos_
#         countPos_ = 0  # 重置计数器
#         print("Button clicked!")
#
#     dpg.add_button(label="Reset Position", callback=button_callback)
#
# # 启动 Dear PyGui 视口
# dpg.setup_dearpygui()
# dpg.show_viewport()

link_name = "left_ankle_roll_link" #"base_link"# 足端 和 base
link_id = m.body(name=link_name).id

with open('robot_state.txt', 'w') as file:
    with mujoco.viewer.launch_passive(m, d) as viewer:
      start = time.time()

      while viewer.is_running() and time.time() - start < 60: # 采集 60s的数据
        step_start = time.time()

        countPos_ += 1
        kkc = 0.5*(np.sin(countPos_ * 0.002 * 2*np.pi) + 1)
        setPos = targetPos * (1 - kkc) + kkc * posDef2

        d.ctrl[:] = targetKp * (setPos - d.qpos[-12:]) + targetKd * (targetVel - d.qvel[-12:])
        mujoco.mj_step(m, d)

        # 获取link的位置姿态 world坐标系
        link_position = d.xpos[link_id]
        link_quaternion = d.xquat[link_id]# wxyz
        link_velocity = np.zeros(6) #  [angular; linear]
        #  获取link的world速度
        mujoco.mj_objectVelocity(m, d, mujoco.mjtObj.mjOBJ_XBODY, link_id, link_velocity, 0)

        # 获取 qpos 和 qvel 的值
        qpos = d.qpos.copy() # 获取 base的xyz 位置世界坐标系，四元数wxyz相对世界系，关节角度
        qvel = d.qvel.copy() # base的xyz速度世界坐标系，base角速度xyz在body坐标系，关节速度


        for value in qpos:
            file.write(f"{value} ")

        for value in qvel:
            file.write(f"{value} ")

        for value in link_position:
            file.write(f"{value} ")

        for value in link_quaternion:
            file.write(f"{value} ")

        for value in link_velocity:
            file.write(f"{value} ")

        file.write("\n")

        # 查看器选项的修改示例：每两秒钟切换一次接触点。
        with viewer.lock():
          viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(d.time % 2)

        # 获取物理状态的更改，应用扰动，从GUI更新选项。
        viewer.sync()

        # 粗略的计时，相对于挂钟会有漂移。
        time_until_next_step = m.opt.timestep - (time.time() - step_start)
        if time_until_next_step > 0:
          time.sleep(time_until_next_step)

        # dpg.render_dearpygui_frame()